import json 
# MODEL_ZOO = ['resnet18','resnet34','resnet50']
# MODEL_ZOO = ['resnet18','resnet34']
# MODEL_BLOCKS = {
#     'resnet18':['layer1.0','layer1.1','layer2.0','layer2.1',
#                 'layer3.0','layer3.1','layer4.0','layer4.1'],
#     'resnet34':['layer1.0','layer1.1','layer1.2',
#                 'layer2.0','layer2.1','layer2.2','layer2.3',
#                 'layer3.0','layer3.1','layer3.2','layer3.3','layer3.4','layer3.5',
#                 'layer4.0','layer4.1','layer4.2'],
#     'resnet50':['layer1.0', 'layer1.1', 'layer1.2',
#                 'layer2.0', 'layer2.1', 'layer2.2', 'layer2.3',
#                 'layer3.0', 'layer3.1', 'layer3.2', 'layer3.3', 'layer3.4', 'layer3.5',
#                 'layer4.0', 'layer4.1', 'layer4.2'],
# }

# MODEL_ZOO = ['resnet18','resnet50','MobileNetV3']
# MODEL_BLOCKS = {       
#     # 'resnet8':[['layer0'],
#     #                     ['layer1.0'],
#     #                     ['layer2.0'],
#     #                     ['layer3.0','avgpool','fc']],
#     'resnet18':[['layer0','layer1.0','layer1.1'],
#                 ['layer2.0','layer2.1'],
#                 ['layer3.0','layer3.1'],
#                 ['layer4.0','layer4.1','avgpool','fc']],
#     'resnet50':[['layer0','layer1.0', 'layer1.1', 'layer1.2','layer2.0', 'layer2.1', 'layer2.2'],
#                 ['layer2.3'],
#                 ['layer3.0', 'layer3.1', 'layer3.2', 'layer3.3', 'layer3.4', 'layer3.5'],
#                 ['layer4.0', 'layer4.1', 'layer4.2','avgpool','fc']],
#     'MobileNetV3':[['conv1','bn1','hs1','layers.0','layers.1','layers.2','layers.3','layers.4',],
#                    ['layers.5',],
#                    ['layers.6','layers.7','layers.8','layers.9','layers.10',],
#                    ['layers.11','layers.12','layers.13','layers.14','layers.15','conv2','bn2','hs2','avgpool','fc']],

# }
# MODEL_ZOO = ['resnet18','resnet50','MobileNetV2']
# MODEL_BLOCKS = {       
#     'resnet18':[['layer0','layer1.0','layer1.1'],
#                 ['layer2.0','layer2.1'],
#                 ['layer3.0','layer3.1'],
#                 ['layer4.0','layer4.1','avgpool','fc']],
#     'resnet50':[['layer0','layer1.0', 'layer1.1', 'layer1.2','layer2.0', 'layer2.1', 'layer2.2'],
#                 ['layer2.3'],
#                 ['layer3.0', 'layer3.1', 'layer3.2', 'layer3.3', 'layer3.4', 'layer3.5'],
#                 ['layer4.0', 'layer4.1', 'layer4.2','avgpool','fc']],
#     'MobileNetV2':[['pre.0','pre.1','pre.2','stage1','stage2.0','stage2.1'],
#                    ['stage3.0','stage3.1','stage3.2'],
#                    ['stage4.0','stage4.1','stage4.2','stage4.3','stage5.0','stage5.1','stage5.2'],
#                    ['stage6.0','stage6.1','stage6.2','stage7','conv1.0','conv1.1','conv1.2','avgpool','conv2']],
# }
####################################################################################
# MODEL_ZOO = ['MobileNetV2','resnet18','vgg16']
# MODEL_BLOCKS = {
#     'MobileNetV2':[['Conv1.0','Conv1.1','Conv1.2','stage1','stage2.0','stage2.1'],
#                 ['stage3.0','stage3.1','stage3.2'],
#                 ['stage4.0','stage4.1','stage4.2','stage4.3','stage5.0','stage5.1','stage5.2'],
#                 ['stage6.0','stage6.1','stage6.2','stage7.0','Conv2.0','Conv2.1','Conv2.2','avgpool','Conv3','flatten']],
#     'resnet18':[['layer0','layer1.0','layer1.1'],
#                 ['layer2.0','layer2.1'],
#                 ['layer3.0','layer3.1'],
#                 ['layer4.0','layer4.1','avgpool','flatten','fc']],
#     'vgg16':[['features1.0','features1.1','features1.2','features1.3','features1.4','features1.5','features1.6'],
#              ['features2.0','features2.1','features2.2','features2.3','features2.4','features2.5','features2.6'],
#              ['features3.0','features3.1','features3.2','features3.3','features3.4','features3.5','features3.6','features3.7','features3.8','features3.9','features4.0','features4.1','features4.2','features4.3','features4.4','features4.5','features4.6','features4.7','features4.8','features4.9','features4.10','features4.11','features4.12','features4.13','features4.14','features4.15','features4.16','features4.17','features4.18','features4.19','features4.20'],
#              ['flatten','classifier.0','classifier.1','classifier.2','classifier.3','classifier.4','classifier.5','classifier.6']],
# }
# # PLANES = [
# #     [[32,256,512],[128,256,512],[256,256,512]],
# #     [[512,256,32],[512,256,128],[512,256,256]]
# # ]
# PLANES = [
#     [[32,128],[128,128],[128,128]],
#     [[128,32],[128,128],[128,128]]
# ]


###################################################################
# PLANES = [
#     [[32,128,256],[128,128,256],[256,256,256]],
#     [[256,128,32],[256,128,128],[256,256,256]]
# ]
# MODEL_ZOO = ['resnet18','resnet34','MobileNetV3','resnet20']
# MODEL_BLOCKS = {       
#     # 'resnet8':[['layer0'],
#     #                     ['layer1.0'],
#     #                     ['layer2.0'],
#     #                     ['layer3.0','avgpool','fc']],
#     'resnet18':[['layer0','layer1.0','layer1.1'],
#                 ['layer2.0','layer2.1'],
#                 ['layer3.0','layer3.1'],
#                 ['layer4.0','layer4.1','avgpool','fc']],
#     'resnet34':[['layer0','layer1.0', 'layer1.1', 'layer1.2','layer2.0', 'layer2.1', 'layer2.2'],
#                 ['layer2.3'],
#                 ['layer3.0', 'layer3.1', 'layer3.2', 'layer3.3', 'layer3.4', 'layer3.5'],
#                 ['layer4.0', 'layer4.1', 'layer4.2','avgpool','fc']],
#     'MobileNetV3':[['conv1','bn1','hs1','layers.0','layers.1','layers.2','layers.3','layers.4',],
#                    ['layers.5',],
#                    ['layers.6','layers.7','layers.8','layers.9','layers.10',],
#                    ['layers.11','layers.12','layers.13','layers.14','layers.15','conv2','bn2','hs2','avgpool','fc']],
#     'MobileNetV2':[['pre','stage1','stage2'],
#                    ['stage3'],
#                    ['stage4','stage5'],
#                    ['stage6','stage7','conv1','','conv2']],
#     'resnet56':[['layer0','layer1.0','layer1.1','layer1.2','layer1.3','layer1.4'],
#                 ['layer1.5'],
#                 ['layer2.0', 'layer2.1', 'layer2.2', 'layer2.3', 'layer2.4', 'layer2.5'],
#                 ['layer3.0', 'layer3.1', 'layer3.2', 'layer3.3', 'layer3.4', 'layer3.5','avgpool','fc']],
# }
# 'MobileNetV2':[[],
#                ]
    # 'regnetY_800mf':[[['layer0','layer1.0','layer1.1'],
    #                 ['layer2.0','layer2.1'],
    #                 ['layer3.0','layer3.1'],
    #                 ['layer4.0','layer4.1','avgpool','fc']]

    #                 ]
#    'resnet8':[['layer0'],
#                         ['layer1.0'],
#                         ['layer2.0'],
#                         ['layer3.0','avgpool','fc']]

# TEMP = [
#     [3,128],[128,128],[128,256],[256,512]
#     [3,64],[64,128],[128,256],[256,512],
#     [3,512],[512,512],[512,1024],[1024,2048],
# ]
# IN_PLANES = [
#     [64,512],[128,512],[256,1024]
# ]
# OUT_PLANES = [
#     [64,512],[128,512],[256,1024]
# ]
# PLANES = [
#     [64,512],[128,512],[256,1024]
# ]
# PLANES = [
#     [[128,512],[512,512]],
#     [[512,128],[512,512]]
# ]
# PLANES = [
#     [[128,256,512],[512,512,512],[32,256,512]],
#     [[512,256,128],[512,512,512],[512,256,32]]
# ]
# PLANES = [
#     [[128,256,512],[512,512,512],[40,160,512],[128,256,512]],
#     [[512,256,128],[512,512,512],[512,160,40],[512,256,128]]
# ]
# PLANES =  [
#     [128,64],[128,128],[256,256]
# ]


MODEL_ZOO = ['cnn','MobileNetV2','resnet18']
MODEL_BLOCKS = {
    'cnn':[['conv1','relu1'],
           ['pool1'],
           ['conv2','relu2','pool2','avgpool'],
           ['flatten','fc1','relu3','fc2','relu4','fc3']
    ],
    'MobileNetV2':[['Conv1.0','Conv1.1','Conv1.2','stage1','stage2.0','stage2.1'],
                ['stage3.0','stage3.1','stage3.2'],
                ['stage4.0','stage4.1','stage4.2','stage4.3','stage5.0','stage5.1','stage5.2'],
                ['stage6.0','stage6.1','stage6.2','stage7.0','Conv2.0','Conv2.1','Conv2.2','avgpool','Conv3','flatten']],
    'resnet18':[['layer0','layer1.0','layer1.1'],
                ['layer2.0','layer2.1'],
                ['layer3.0','layer3.1'],
                ['layer4.0','layer4.1','avgpool','flatten','fc']],
}
PLANES = [
    [[16,128],[32,128],[128,128]],
    [[128,16],[128,32],[128,128]]
]

# MODEL_ZOO = ['cnn','MobileNetV2','resnet18']
# MODEL_BLOCKS = {
#     'cnn':[['reshape.0','reshape.1','reshape.2','reshape.3','conv1','relu1'],
#            ['pool1'],
#            ['conv2','relu2','pool2','avgpool'],
#            ['flatten','fc1','relu3','fc2','relu4','fc3']
#     ],
#     'MobileNetV2':[['reshape.0','reshape.1','reshape.2','reshape.3','Conv1.0','Conv1.1','Conv1.2','stage1','stage2.0','stage2.1'],
#                 ['stage3.0','stage3.1','stage3.2'],
#                 ['stage4.0','stage4.1','stage4.2','stage4.3','stage5.0','stage5.1','stage5.2'],
#                 ['stage6.0','stage6.1','stage6.2','stage7.0','Conv2.0','Conv2.1','Conv2.2','avgpool','Conv3','flatten']],
#     'resnet18':[['reshape.0','reshape.1','reshape.2','reshape.3','layer0','layer1.0','layer1.1'],
#                 ['layer2.0','layer2.1'],
#                 ['layer3.0','layer3.1'],
#                 ['layer4.0','layer4.1','avgpool','flatten','fc']],
# }
# PLANES = [
#     [[16,128],[32,128],[128,128]],
#     [[128,16],[128,32],[128,128]]
# ]

